package com.lambo.los.kits.war;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Method;
import java.net.*;
import java.util.*;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.Manifest;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;

/**
 * war包处理器.
 *
 * @author 林小宝
 * @createTime 2014年8月10日 上午1:01:48
 */
public class Handler extends URLStreamHandler {
    private static final Map<String, byte[]> cache = new HashMap<>(1024);
    private static final String PROTOCOL;
    private static final String PROTOCOL_SPLIT = "!/";
    static {
        String JAVA_PROTOCOL_HANDLER = "java.protocol.handler.pkgs";
        String handlerPackage = System.getProperty(JAVA_PROTOCOL_HANDLER);
        if (handlerPackage == null) handlerPackage = "";
        if (handlerPackage.length() > 0) handlerPackage = "|" + handlerPackage;
        String packageName = Handler.class.getName();
        packageName = packageName.substring(0, packageName.length() - 8);
        PROTOCOL = packageName.substring(packageName.lastIndexOf(".") + 1);
        packageName = packageName.substring(0, packageName.length() - PROTOCOL.length() - 1);
        handlerPackage = packageName + handlerPackage;
        System.setProperty(JAVA_PROTOCOL_HANDLER, handlerPackage);
    }

    protected URLConnection openConnection(final URL u) throws IOException {
        return new URLConnection(u) {
            @Override
            public void connect() throws IOException {
            }

            public InputStream getInputStream() throws IOException {
                return new ByteArrayInputStream(cache.get(u.toString()));
            }
        };
    }

    @Override
    protected int hashCode(URL u) {
        if(null == u) {
            return 0;
        }
        return u.toString().hashCode();
    }

    public static void main(String[] args) throws Exception {
        System.setProperty("system.start", String.valueOf(System.currentTimeMillis()));
        String lib = "WEB-INF/lib";
        URL resource = Handler.class.getClassLoader().getResource(lib);
        if (null == resource) {
            throw new RuntimeException("only support has WEB-INF/lib");
        }
        JarURLConnection urlConnection = (JarURLConnection) resource.openConnection();
        JarFile jarFile = urlConnection.getJarFile();

        Manifest manifest = jarFile.getManifest();
        Attributes mainAttributes = manifest.getMainAttributes();
        Enumeration<JarEntry> entries = jarFile.entries();
        URLClassLoader contextClassLoader = new LocalURLClassLoader(new URL[0], Thread.currentThread().getContextClassLoader());
        Thread.currentThread().setContextClassLoader(contextClassLoader);
        while (entries.hasMoreElements()) {
            JarEntry jarEntry = entries.nextElement();
            if (jarEntry.getName().startsWith(lib) && jarEntry.getName().toLowerCase().endsWith(".jar")) {
                byte[] data = readBytes(jarFile.getInputStream(jarEntry));
                ZipInputStream zipInputStream = new ZipInputStream(new ByteArrayInputStream(data));
                while (true) {
                    ZipEntry nextEntry = zipInputStream.getNextEntry();
                    if (null == nextEntry) {
                        break;
                    }
                    String name = nextEntry.getName();
                    if (name.endsWith("/")) {
                        cache.put("jar:" + PROTOCOL + "://" + jarEntry.getName() + PROTOCOL_SPLIT + name, data);
                        continue;
                    }
                    cache.put(PROTOCOL + "://" + jarEntry.getName() + PROTOCOL_SPLIT + name, readBytes(zipInputStream));
                }
            }
        }
        jarFile.close();

        String warMainClass = mainAttributes.getValue("warMainClass");
        Class<?> aClass = contextClassLoader.loadClass(warMainClass);
        Method main = aClass.getMethod("main", String[].class);
        main.invoke(aClass, (Object) args);
    }

    private static byte[] readBytes(InputStream inputStream) throws IOException {
        ByteArrayOutputStream tmp = new ByteArrayOutputStream();
        int n;
        byte[] data = new byte[102400];
        while ((n = inputStream.read(data)) > -1) {
            tmp.write(data, 0, n);
        }
        return tmp.toByteArray();
    }

    public static class LocalURLClassLoader extends URLClassLoader {
        public LocalURLClassLoader(URL[] urls, ClassLoader contextClassLoader) {
            super(urls, contextClassLoader);
        }

        @Override
        public URL findResource(String name) {
            String cacheName = PROTOCOL_SPLIT + name;
            for (Map.Entry<String, byte[]> entry : cache.entrySet()) {
                if (entry.getKey().endsWith(cacheName)) {
                    try {
                        return new URL(entry.getKey());
                    } catch (MalformedURLException e) {
                        e.printStackTrace();
                    }
                }
            }
            return null;
        }

        @Override
        public Enumeration<URL> findResources(String name) throws IOException {
            Vector<URL> result = new Vector<>();
            String cacheName = PROTOCOL_SPLIT + name;
            try {
                for (Map.Entry<String, byte[]> entry : cache.entrySet()) {
                    if (entry.getKey().endsWith(cacheName)) {
                        result.add(new URL(entry.getKey()));
                    }
                }
            } catch (MalformedURLException e) {
                e.printStackTrace();
            }
            return result.elements();
        }

        @Override
        public Class<?> findClass(String name) throws ClassNotFoundException {
            String cacheName = PROTOCOL_SPLIT + name.replace(".", "/") + ".class";
            for (Map.Entry<String, byte[]> entry : cache.entrySet()) {
                if (entry.getKey().endsWith(cacheName)) {
                    return super.defineClass(name, entry.getValue(), 0, entry.getValue().length);
                }
            }
            return super.loadClass(name);
        }
    }
}
